"""
author: liuxu
date: 2024.12.3
description: 1、hiprof --leak-check资源泄漏检测模块
            2、可与valgrind工具对比结果（valgrind --leak-check=full --show-leak-kinds=all ./my_program）
            3、支持检测的函数详见：hip_prof/libs/leakcheckmodule/include/hip_profile_leakcheck.h
"""
import re
import allure
import pytest
import pathlib

from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd
from common.common_func import convert_to_bytes
from base_public.testbase import BasePublic


@pytest.mark.high
@allure.feature('内存泄漏模块测试')
class TestLeakCheck(BasePublic):
    def test_matrix_leak(self, test_env):
        """
        测试矩阵转置程序
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="hipSimple/matrix.cpp", des_exe_name="matrix-O0-g",
                                            compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)


    def test_cpu_memory_leak(self, test_env):
        """
        测试cpu内存泄漏
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/cpu_memory_test.cpp",
                                            des_exe_name="cpu_memory_test-O0-g",
                                            compile_paras="-O0 -g")

        # 执行内存泄漏测试
        malloc_size = 1024 * 1024  # 1MB
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path} --size {malloc_size}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        # 检测输出的内存泄漏函数和大小是否正确
        expect_func = ["new", "malloc", "calloc", "realloc", "mmap"]  # 程序实际泄漏的内存
        pattern = r"fun_name:(\w+).*?total_size:([0-9.]+[KMG]?B)"
        matches = re.findall(pattern, output)
        real_func = set()
        for func, size in matches:
            # print(f"Function: {func}, Total Size: {size}")
            if func in expect_func and convert_to_bytes(size) == malloc_size:
                real_func.add(func)
        assert real_func == set(expect_func)

    def test_cpu_memory_free(self, test_env):
        """
        测试cpu内存申请后，均已释放
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/cpu_memory_test.cpp",
                                            des_exe_name="cpu_memory_test-O0-g",
                                            compile_paras="-O0 -g")

        # 执行内存泄漏测试
        malloc_size = 1024 * 1024
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path} --size {malloc_size} --free", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        # 检测输出的内存泄漏函数和大小
        expect_func = ["new", "malloc", "calloc", "realloc", "mmap"]  # 程序实际泄漏的内存大小
        pattern = r"fun_name:(\w+).*?total_size:([0-9.]+[KMG]?B)"
        matches = re.findall(pattern, output)
        for func, size in matches:
            # print(f"Function: {func}, Total Size: {size}")
            if func in expect_func and convert_to_bytes(size) == malloc_size:
                assert False, f"{func} did not meet expectations"

    def test_gpu_memory_leak(self, test_env):
        """
        测试gpu内存泄漏
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/gpu_memory_test.cpp",
                                            des_exe_name="gpu_memory_test-O0-g",  # 用例包含多个hip内存操作接口
                                            compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check --show-limit 10000 {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

        # 程序调用了hipMalloc hipHostMalloc hipMallocManaged hipMalloc3D等接口，其中hipMalloc调用3次，其余均调用1次，都没释放内存
        leak_matches = self.get_leak_func_info_form_output(output)
        mylogger.info(f"hip api memory leak check info:{leak_matches}")
        leak_func_list = [m[0] for m in leak_matches]
        assert leak_func_list.count("hipMalloc") == 3
        assert leak_func_list.count("hipHostMalloc") == 1
        assert leak_func_list.count("hipMallocManaged") == 1
        assert leak_func_list.count("hipMalloc3D") == 1

    def test_gpu_memory_free(self, test_env):
        """
        测试gpu内存申请后已释放
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/gpu_memory_test.cpp",
                                            des_exe_name="gpu_memory_test-O0-g",  # 用例包含多个hip内存操作接口
                                            compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check --show-limit 10000 {exe_path} --free", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

        # 程序调用了hipMalloc hipHostMalloc hipMallocManaged hipMalloc3D等接口，都已经释放内存
        p = re.compile(r"HIP_PROF:\s+fun_name:(hip\w+)\s+times:(\d+)\s+total_size:(\d+(?:\.\d+)?\w+)\s+device_type:(\w+)",
                       re.MULTILINE)
        matches = re.findall(p, output.split("leak check start")[-1])
        mylogger.info(f"hip api memory leak check info:{matches}")
        assert len(matches) == 0

    def test_stream_and_event_leak(self, test_env):
        """
        测试流和事件泄漏
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/stream_event_leak_test.cpp")

        # 执行内存泄漏测试
        # hipStreamCreate,hipStreamCreateWithFlags,hipStreamCreateWithPriority,hipEventCreate,hipEventCreateWithFlags泄漏
        ret, output = exec_cmd(f"hipprof --leak-check --show-limit 100 {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

        leak_matches = self.get_leak_func_info_form_output(output)
        leak_func_list = [m[0] for m in leak_matches]
        assert leak_func_list.count("hipStreamCreate") == 1
        assert leak_func_list.count("hipStreamCreateWithFlags") == 1
        assert leak_func_list.count("hipStreamCreateWithPriority") == 1
        assert leak_func_list.count("hipEventCreate") == 1
        assert leak_func_list.count("hipEventCreateWithFlags") == 1

    @allure.link("http://hpczentao.sugon.com/story-view-6924-1-project-485.html", name="禅道需求链接")
    def test_virtual_memory_leak(self, test_env):
        """
        测试虚拟内存泄漏
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/virtual_memory_leak_test.cpp")

        # 执行内存泄漏测试
        size_bytes = 111*1024
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path} --size={size_bytes}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        # 检测获取内存泄漏信息函数和大小
        for l in self.get_leak_func_info_form_output(output):
            if l[0] == "mmap":
                assert convert_to_bytes(l[2]) == size_bytes
                return
        else:
            pytest.fail("mmap memory leak information not found")

    @pytest.mark.parametrize("input_limit", (1, 10, 30))
    def test_leak_check_show_limit(self, input_limit, test_env):
        """
        测试限制泄漏函数堆栈的数量
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/gpu_memory_test.cpp")

        # 执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check --show-limit {input_limit} {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        leak_matches = self.get_leak_func_info_form_output(output)
        assert len(leak_matches) == input_limit

    def test_leak_check_minsize(self, test_env):
        """
        测试内存泄漏检测大小低于默认值
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/cpu_memory_test.cpp")

        # 执行内存泄漏测试
        min_size = 64  # leak-check默认最小检测1024B
        ret, output = exec_cmd(
            f"hipprof --leak-check --check-minsize {min_size} --show-limit 10000 {exe_path} --size {min_size}",
            cwd=test_env,
            timeout=60)
        #TODO 偶现输出数据量过大，subprocess获取日志失败，待解决
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        expect_func = ["new", "malloc", "calloc", "realloc", "mmap"]
        for func in expect_func:
            assert f"fun_name:{func}" in output.split("leak check start")[1]

    def test_leak_check_large_size(self, test_env):
        """
        测试内存泄漏检测大小高于默认值
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="hipSimple/matrix.cpp")

        # 执行内存泄漏测试
        size = 20 * 1024 *1024  # 设置20MB
        ret, output = exec_cmd(f"hipprof --leak-check --check-minsize {size} --show-limit 10000 {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        # 只输出泄漏大于size的堆栈
        for (fun_name, times, total_size, device_type) in self.get_leak_func_info_form_output(output):
            assert convert_to_bytes(total_size) >= size

    def test_gen_leak_info_from_db(self, test_env):
        """
        测试--time-range和--db功能
        :param test_env:
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="hipSimple/matrix.cpp")

        # 执行内存泄漏测试
        # step1: 运行程序生成db，程序运行小于1s
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)
        ori_leak = self.get_leak_func_info_form_output(output)
        # 获取生成的db文件名
        pid = self.get_prof_pid(output)
        db_name = f"hip-prof-{pid}.db"

        # step2: 通过db生成leak-check数据，指定--time-range时间范围包括程序运行时间
        ret2, output2 = exec_cmd(f"hipprof --leak-check --time-range 0:1 --db {db_name}", cwd=test_env)
        assert ret2 is True
        export_leak2 = self.get_leak_func_info_form_output(output2)
        assert export_leak2 == ori_leak

        # step3：通过db生成leak-check数据，指定--time-range时间范围超出程序运行时间，导出数据应该为空
        ret3, output3 = exec_cmd(f"hipprof --leak-check --time-range 10:20 --db {db_name}", cwd=test_env)
        assert ret3 is True
        export_leak3 = self.get_leak_func_info_form_output(output3)
        assert len(export_leak3) == 0

    @pytest.mark.high
    def test_record_memfreeinfo(self, test_env):
        """
        测试记录所有的内存申请信息(包括已释放的内存信息)
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/record-memfreeinfo_test.cpp",
                                            des_exe_name="record-memfreeinfo_test-O0-g", compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret0, output0 = exec_cmd(f"hipprof --leak-check --record-memfreeinfo {exe_path}", cwd=test_env)
        assert ret0 is True
        assert self.check_leakcheck_is_success(output0)

        # 获取db文件
        prof_pid = self.get_prof_pid(output0)

        # 从db导出泄漏堆栈，包括未泄漏数据
        ret1, output1 = exec_cmd(f"hipprof --time-range 1:2 --db hip-prof-{prof_pid}.db", cwd=test_env)
        assert ret1 is True
        leak_func_info_l = self.get_leak_func_info_form_output(output1)
        assert len(leak_func_info_l) == 2 and ["hipMalloc" == func_info[0] for func_info in leak_func_info_l]
        # TODO 观察内存曲线图

    @pytest.mark.medium
    def test_multiproc_multidevice(self, test_env):
        """
        测试多进程多卡用例
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="multiproc_case/multiprocess_multithread_multidevice_test.cpp",
                                            des_exe_name="multiprocess_multithread_multidevice_test-O0-g", compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check {exe_path}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

    _scenarios = [
        (
            "leak",  # 场景名称
            "--leak",  # 命令参数
            "All functions that have not released resources should not be detected"  # 预期结果
        ),
        (
            "noleak",  # 场景名称
            "",  # 命令参数（空）
            "No function that has released resources should be detected"  # 预期结果
        )
    ]
    @pytest.mark.parametrize("test_scenario, exe_param, expect_desc", _scenarios)
    def test_leakcheck_2510_added_api(self, test_env, test_scenario, exe_param, expect_desc):
        """
        测试25.10需求新增支持的api
        :param test_env:
        :return:
        """
        # 2510leakcheck新增支持的api
        expect_api_s = {'hipGraphCreate', 'hipArrayCreate', 'hipArray3DCreate', 'hipCtxCreate',
                        'hipMipmappedArrayCreate', 'hipTexRefCreate', 'hipUserObjectCreate', 'hipMemCreate',
                        'hipCreateTextureObject', 'hipGreenCtxCreate', 'hipMemPoolCreate', 'hipCreateSurfaceObject',
                        'hipSurfObjectCreate'}
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="leak_check/hip_resource_leak_test_2510.cpp",
                                            compile_paras="-O0 -g -lm")

        # 在未释放资源场景，执行内存泄漏测试
        ret, output = exec_cmd(f"hipprof --leak-check --show-limit 1000 {exe_path} {exe_param}", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

        # 获取生成的内存泄漏信息
        all_leakcheck_info_l = self.get_leak_func_info_form_output(output)
        test_leakcheck_api_name_s = set([info[0] for info in all_leakcheck_info_l])

        if test_scenario == "leak":
            # 泄漏场景：所有目标API必须全部被检测到
            missing_expected_apis = expect_api_s - test_leakcheck_api_name_s  # 差集：expect_api_s中未被检测出的api
            assert not missing_expected_apis, \
                f"The following API has not been captured by the leakcheck tool: {sorted(missing_expected_apis)}"
        else:
            # 未泄漏场景：任一测试API都不应被检测到
            unexpected_apis = expect_api_s & test_leakcheck_api_name_s  # 交集
            assert not unexpected_apis, \
                f"The following API has been captured by the leakcheck tool: {sorted(unexpected_apis)}"
            # TODO测试结果不稳定，偶尔会误报hipMipmappedArrayCreate、hipArrayCreate
        # hipcc hip_memory_leak_test_fixed.cpp -O0 -g -lm -o hip_memory_leak_test_fixed
        # hipprof --leak-check --show-limit 1000 --check-minsize 1 ./a --leak | grep "func_name:hip"

    def test_graph_leak(self, test_env):
        """
        测试图相关的泄漏
        :param test_env:
        :return:
        """
        # 测试的泄漏函数
        # expect_api_s = {'hipGraphCreate', 'hipGraphAddMemAllocNode', 'hipGraphRetainUserObject', 'hipUserObjectCreate'}
        expect_api_s = {'hipGraphCreate', 'hipUserObjectCreate'}

        # 编译待测用例
        # exe_path = self.compile_source_code(src_file_name="graph_case/ahipGraphTest_height.cpp")
        exe_path = self.compile_source_code(src_file_name="leak_check/hipGraph_leaktest.cpp", compile_paras="-O0 -g")

        # 执行内存泄漏测试
        ret, output = exec_cmd(
            f"hipprof --leak-check --check-minsize {10 * 1024 * 1024} --show-limit 100 {exe_path} --leak", cwd=test_env)
        assert ret is True
        assert self.check_leakcheck_is_success(output)

        # 获取生成的内存泄漏信息
        all_leakcheck_info_l = self.get_leak_func_info_form_output(output)
        test_leakcheck_api_name_s = set([info[0] for info in all_leakcheck_info_l])

        # 泄漏场景：所有目标API必须全部被检测到
        missing_expected_apis = expect_api_s - test_leakcheck_api_name_s  # 差集：获取未被检测出的接口
        if missing_expected_apis:
            # 通过hipprof --hip-trace检查api均正常执行
            exec_cmd(f"hipprof --hip-trace {exe_path} --leak", cwd=test_env)

            assert False, \
                f"The following API has not been captured by the leakcheck tool: {sorted(missing_expected_apis)}"
        # 'hipGraphAddMemAllocNode', 'hipGraphRetainUserObject'这两个api测不出来


@pytest.mark.smoke
@pytest.mark.high
@allure.feature('内存使用曲线测试')
class TestMemoryUsage(BasePublic):
    def test_memory_usage(self, test_env):
        """
        测试cpu/gpu内存使用曲线
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="memory_usage/multidevice_memory_usage.cpp")

        # 执行内存泄漏测试
        # --memory-usage只会生成hip-prof-{pid}.db和内存使用曲线mem.html文件
        ret, output = exec_cmd(f"hipprof --memory-usage {exe_path}", cwd=test_env, timeout=60)
        assert ret is True
        pid = self.get_prof_pid(output)
        matches = pathlib.Path(test_env).rglob(f"result_{pid}.mem*.html")
        assert len(list(matches)) >= 1
        # TODO 观察内存曲线图



